{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "23842a6f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import os\n",
    "import sys\n",
    "from mqbench.convert_deploy import convert_deploy\n",
    "from mqbench.prepare_by_platform import prepare_by_platform, BackendType\n",
    "from mqbench.utils.state import enable_calibration, enable_quantization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "75ec1646",
   "metadata": {},
   "outputs": [],
   "source": [
    "import mqbench\n",
    "from mqbench.convert_deploy import convert_deploy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "840f6145",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torchvision.models as models\n",
    "import torchvision.datasets as datasets\n",
    "import torchvision.transforms as transforms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "148e995c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb01e579",
   "metadata": {},
   "outputs": [],
   "source": [
    "resnet18 = models.resnet18(pretrained=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24f5b3fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "def test(model, device, test_loader):\n",
    "    old_training_state = model.training\n",
    "    model.eval()\n",
    "    test_loss = 0\n",
    "    correct = 0\n",
    "    lossLayer = torch.nn.CrossEntropyLoss(reduction='sum')\n",
    "    for data, target in tqdm.tqdm(test_loader):\n",
    "        data, target = data.to(device), target.to(device)\n",
    "        with torch.no_grad():\n",
    "            output = model(data)\n",
    "        test_loss += lossLayer(output, target).item()\n",
    "        pred = output.argmax(dim=1, keepdim=True)\n",
    "        correct += pred.eq(target.view_as(pred)).sum().item()\n",
    "    \n",
    "    test_loss /= len(test_loader.dataset)\n",
    "    model.train(old_training_state)\n",
    "    print('\\nTest set: Average loss: {:.4f}, Accuracy: {:.3f}%\\n'.format(\n",
    "        test_loss, 100. * correct / len(test_loader.dataset)\n",
    "    ))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ffab05f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# change valdir to your imagenet dataset validation directory\n",
    "valdir = 'yourdisk/ImageNetDataset/val/'\n",
    "normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.225])\n",
    "\n",
    "val_loader = torch.utils.data.DataLoader(\n",
    "        datasets.ImageFolder(valdir, transforms.Compose([\n",
    "            transforms.Resize(256),\n",
    "            transforms.CenterCrop(224),\n",
    "            transforms.ToTensor(),\n",
    "            normalize,\n",
    "        ])),batch_size=128, shuffle=False, num_workers=24, pin_memory=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd656f97",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device('cuda')\n",
    "resnet18.to(device)\n",
    "# 69.758, consistent with reported at https://pytorch.org/vision/stable/models.html\n",
    "test(resnet18, device, val_loader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93f93ef5",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = resnet18.cpu().train()\n",
    "model_mqbench = prepare_by_platform(model, BackendType.OPENVINO)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6cdceda",
   "metadata": {},
   "outputs": [],
   "source": [
    "enable_calibration(model_mqbench)\n",
    "\n",
    "model_mqbench.to(device)\n",
    "model_mqbench.eval()\n",
    "# calibration loop\n",
    "for data, target in tqdm.tqdm(val_loader):\n",
    "        data, target = data.to(device), target.to(device)\n",
    "        with torch.no_grad():\n",
    "            model_mqbench(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55b06a5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "enable_quantization(model_mqbench)\n",
    "model_mqbench.train()\n",
    "# QAT loop ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "202fb4e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "test(model_mqbench, device, val_loader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3a5bd73",
   "metadata": {},
   "outputs": [],
   "source": [
    "input_shape={'data': [32, 3, 224, 224]}\n",
    "convert_deploy(model_mqbench.eval(), BackendType.OPENVINO, input_shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "83f96535",
   "metadata": {},
   "source": [
    "# test openvino"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05b1c0a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "!mo --input_model mqbench_qmodel_deploy_model.onnx "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4533dd4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging as log\n",
    "import numpy as np\n",
    "from openvino.inference_engine import IECore"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd13db8a",
   "metadata": {},
   "outputs": [],
   "source": [
    "log.basicConfig(format=\"[ %(levelname)s ] %(message)s\", level=log.ERROR, stream=sys.stdout)\n",
    "log.info(\"Creating Inference Engine...\")\n",
    "ie = IECore()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04cf7ec9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Read IR\n",
    "log.info(\"Loading network\")\n",
    "net = ie.read_network(\"./mqbench_qmodel_deploy_model.xml\")\n",
    "\n",
    "img_info_input_blob = None\n",
    "feed_dict = {}\n",
    "input_blob = \"input\"\n",
    "for blob_name in net.input_info:\n",
    "    if len(net.input_info[blob_name].input_data.shape) == 4:\n",
    "        input_blob = blob_name\n",
    "    elif len(net.input_info[blob_name].input_data.shape) == 2:\n",
    "        img_info_input_blob = blob_name\n",
    "    else:\n",
    "        raise RuntimeError(\"Unsupported {}D input layer '{}'. Only 2D and 4D input layers are supported\"\n",
    "                           .format(len(net.input_info[blob_name].input_data.shape), blob_name))\n",
    "\n",
    "log.info(\"Loading IR to the plugin...\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b1f4c66",
   "metadata": {},
   "outputs": [],
   "source": [
    "exec_net = ie.load_network(network=net, num_requests=4, device_name=\"CPU\")\n",
    "n, c, h, w = net.input_info[input_blob].input_data.shape\n",
    "if img_info_input_blob:\n",
    "    feed_dict[img_info_input_blob] = [h, w, 1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4f46b91",
   "metadata": {},
   "outputs": [],
   "source": [
    "n, c, h, w"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20ff84cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn.functional as F"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3720657",
   "metadata": {},
   "outputs": [],
   "source": [
    "val_loader = torch.utils.data.DataLoader(\n",
    "        datasets.ImageFolder(valdir, transforms.Compose([\n",
    "            transforms.Resize(256),\n",
    "            transforms.CenterCrop(224),\n",
    "            transforms.ToTensor(),\n",
    "            normalize,\n",
    "        ])),batch_size=32, shuffle=False, num_workers=24, pin_memory=True, drop_last=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d03f827",
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_openvino(exec_net, test_loader, model):\n",
    "    \n",
    "    global feed_dict\n",
    "    global input_blob\n",
    "    cur_request_id = 0\n",
    "    \n",
    "    test_loss = 0\n",
    "    correct = 0\n",
    "    \n",
    "    def infer(img):\n",
    "        cur_request_id = 0\n",
    "        feed_dict[input_blob] = img\n",
    "        exec_net.start_async(request_id=cur_request_id, inputs=feed_dict)\n",
    "        while exec_net.requests[cur_request_id].wait(-1) != 0:\n",
    "            pass\n",
    "        outs = exec_net.requests[cur_request_id].output_blobs\n",
    "        k = tuple(outs.keys())[0]\n",
    "        return outs[k].buffer\n",
    "    \n",
    "    lossLayer = torch.nn.CrossEntropyLoss(reduction='sum')\n",
    "    device = torch.device('cuda')\n",
    "    old_training_state = model.training\n",
    "    model.to(device)\n",
    "    model.eval()\n",
    "    cosine_distance = 0\n",
    "    for data, target in tqdm.tqdm(test_loader):\n",
    "        output = infer(data.numpy())\n",
    "        output = torch.as_tensor(output)\n",
    "        with torch.no_grad():\n",
    "            output_torch = model(data.to(device)).cpu()\n",
    "        cosine_distance += (F.cosine_similarity(output, output_torch)).sum()\n",
    "        test_loss += lossLayer(output, target).item()\n",
    "        pred = output.argmax(dim=1, keepdim=True)\n",
    "        correct += pred.eq(target.view_as(pred)).sum().item()\n",
    "    cosine_distance /= len(test_loader.dataset)\n",
    "    test_loss /= len(test_loader.dataset)\n",
    "    model.train(old_training_state)\n",
    "    print('\\nTest set: Average loss: {:.4f}, Accuracy: {:.3f%}, cosine_distance:{:.4f}\\n'.format(\n",
    "        test_loss, 100. * correct / len(test_loader.dataset), cosine_distance\n",
    "    ))\n",
    "    return cosine_distance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33cede37",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_openvino(exec_net, val_loader, model_mqbench)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e543b8c",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.onnx.export(resnet18.cpu(), torch.rand(1,3,224,224), 'resnet18.onnx', do_constant_folding=True,\n",
    "                      input_names=[\"input\"], output_names=[\"output\"], opset_version=11)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0067f15c",
   "metadata": {},
   "outputs": [],
   "source": [
    "!mo --input_model resnet18.onnx "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9fad2840",
   "metadata": {},
   "outputs": [],
   "source": [
    "input_shape={'data': [1, 3, 224, 224]}\n",
    "convert_deploy(model_mqbench.cpu().eval(), BackendType.OPENVINO, input_shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45e7f5f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "!mo --input_model resnet18.onnx "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6364bcd3",
   "metadata": {},
   "outputs": [],
   "source": [
    "!mo --input_model mqbench_qmodel_deploy_model.onnx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1967c9d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "!benchmark_app -m resnet18.xml -nstreams 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c4a765a",
   "metadata": {},
   "outputs": [],
   "source": [
    "!benchmark_app -m mqbench_qmodel_deploy_model.xml  -nstreams 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01ea3ea7",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
